import argparse
import re

import torch

from models.SENet.SEUnet import SimSEUnet
from utils.train.trainUtils import pretrain, makedirs, snapshotargs, fit_one_epoch2


def main(args):
    makedirs(args)
    device = torch.device("cpu" if not torch.cuda.is_available() else args.device)

    net = SimSEUnet(in_channels=args.inputs_size[-1],
                    out_channels=args.image_size[-1])

    if args.pretrained:
        pretrain(model=net,
                 model_path=args.pretrained_path)
        match = re.search("Epoch(.*)-mean_dsc(......).pth", args.pretrained_path)
        args.init_epoch = int(match.group(1)) - 1
        args.best_validation_dsc = float(match.group(2))

    snapshotargs(args)

    fit_one_epoch2(net=net, device=device, args=args)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Training SE-U-Net model for segmentation of RVSC"
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=4,
        help="input batch size for training (default: 2)",
    )
    parser.add_argument(
        "--init_epoch",
        type=int,
        default=0,
        help="the initial generation (default: 0)",
    )
    parser.add_argument(
        "--epochs",
        type=int,
        default=300,
        help="number of epochs to train (default: 300)",
    )
    parser.add_argument(
        "--lr",
        type=float,
        default=1e-4,
        help="initial learning rate (default: 1e-4)",
    )
    parser.add_argument(
        "--log_dir",
        type=str,
        default="./checkpoint_SEUnet_1",
        help="folder to save weights"
    )
    parser.add_argument(
        "--pretrained",
        type=bool,
        default=False,
        help="whether to use backbone network pre training weights"
    )
    parser.add_argument(
        "--pretrained_path",
        type=str,
        default=None,
        help="file to continue"
    )
    parser.add_argument(
        "--mask_type",
        type=str,
        default='both',
        help="Number of categories +1"
    )
    parser.add_argument(
        "--inputs_size",
        type=list,
        default=[256, 256, 1],
        help="the size of the input image"
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        help="device for training (default: cuda:0)",
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=4,
        help="number of workers for data loading (default: 4)",
    )
    parser.add_argument(
        "--logs", type=str, default="./logs_seunet", help="folder to save logs"
    )
    parser.add_argument(
        "--trainDataset_path",
        type=str,
        default="../../../data/RVSC2012/TrainingSet",
        help="root folder of trainDataset"
    )
    parser.add_argument(
        "--testDataset_path1",
        type=str,
        default="../../../data/RVSC2012/Test1Set",
        help="root folder of testDataset_path1"
    )
    parser.add_argument(
        "--testDataset_path2",
        type=str,
        default="../../../data/RVSC2012/Test2Set",
        help="root folder of testDataset_path2"
    )
    parser.add_argument(
        "--image_size",
        type=int,
        default=[256, 256, 1],
        help="target input image size (default: 256)",
    )
    parser.add_argument(
        "--best_validation_dsc",
        type=int,
        default=0,
        help="best_validation_dsc (default: 0)",
    )
    args = parser.parse_args()
    main(args)
